# -*- coding: utf-8 -*-

from pylab import *
from scipy import signal

def fwth(x):
    n = int(log2(x.size))
    tx = array(x) * 1.0
    
    for i in xrange(1, n + 1):
        ttx = array(tx)
        jMax = 1 << (n - i)
        for j in xrange(0, jMax):
            l0 = 2 * j
            l1 = l0 + 1
            r0 = j
            r1 = r0 + jMax
            ttx[r0] = (tx[l0] + tx[l1]) / 2
            ttx[r1] = (tx[l0] - tx[l1]) / 2
        tx = array(ttx)
    return tx

def ifwth(x):
    n = int(log2(x.size))
    tx = array(x) * 1.0
    
    for i in xrange(n, 0, -1):
        ttx = array(tx)
        jMax = 1 << (n - i)
        for j in xrange(jMax - 1, -1, -1):
            l0 = j
            l1 = l0 + jMax
            r0 = 2 * j
            r1 = r0 + 1
            ttx[r0] = tx[l0] + tx[l1]
            ttx[r1] = tx[l0] - tx[l1]
        tx = array(ttx)
    return tx

if __name__ == '__main__':
    t = linspace(0, 7, 8)
    sig = signal.sawtooth(t, 0.8)
    t = linspace(0, 8 * 512 - 1, 8 * 512)
    sig = repeat(sig, 512)
    sig += randn(len(sig))
    
    X = fwth(sig)
    figure()
    subplot(3, 1, 1)
    plot(t, sig)
    subplot(3, 1, 2)
    plot(t, X)
    X[2048:len(X)] = 0
    subplot(3, 1, 3)
    plot(t, sig, t, ifwth(X), 'r')
    show()
